from collections.abc import Generator
from typing import Any
import re
import pymysql
from pymysql.connections import Connection
from pymysql.cursors import DictCursor

from dify_plugin import Tool
from dify_plugin.entities.tool import ToolInvokeMessage

class ScarecrowMysqlTool(Tool):
    """一个简单的 MYSQL DDL RUN 执行工具。"""

    def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage]:
        """
        执行MySQL查询的工具实现。
        
        Args:
            tool_parameters: 包含查询SQL的参数字典
                - query_sql: 要执行的SQL查询语句
        
        Yields:
            ToolInvokeMessage: 包含查询结果的消息
        
        Raises:
            Exception: 当配置缺失、SQL非法或执行出错时抛出异常
        """
        mysql_connection: Connection = None
        try:
            # 1. 从运行时获取凭证
            credentials = self.runtime.credentials
            mysql_host = credentials.get("mysql_host")
            mysql_port = int(credentials.get("mysql_port", 3306))
            mysql_user = credentials.get("mysql_user_name")
            mysql_password = credentials.get("mysql_user_pwd")
            mysql_database = credentials.get("mysql_database_name")
            
            # 验证必要的凭证
            required_credentials = ["mysql_host", "mysql_user_name", "mysql_user_pwd", "mysql_database_name"]
            missing_credentials = [cred for cred in required_credentials if not credentials.get(cred)]
            if missing_credentials:
                raise Exception(f"缺少必要的配置项：{', '.join(missing_credentials)}")

            # 2. 获取并验证SQL查询
            query_sql = tool_parameters.get("query_sql", "").strip()
            if not query_sql:
                raise Exception("查询SQL不能为空")

            # 3. 验证SQL是否为查询语句
            if not self._is_select_query(query_sql):
                raise Exception("只允许执行SELECT查询语句")

            # 4. 建立数据库连接
            mysql_connection = pymysql.connect(
                host=mysql_host,
                port=mysql_port,
                user=mysql_user,
                password=mysql_password,
                database=mysql_database,
                charset='utf8mb4',
                cursorclass=DictCursor  # 使用字典游标返回结果
            )

            # 5. 执行查询
            with mysql_connection.cursor() as cursor:
                cursor.execute(query_sql)
                results = cursor.fetchall()
                
                # 6. 返回结果
                yield self.create_json_message({
                    "success": True,
                    "data": results,
                    "total": len(results),
                    "message": "查询执行成功"
                })

        except Exception as e:
            # 发生错误时返回错误信息
            yield self.create_json_message({
                "success": False,
                "data": None,
                "message": f"查询执行失败：{str(e)}"
            })

        finally:
            # 确保连接被正确关闭
            if mysql_connection:
                mysql_connection.close()

    def _is_select_query(self, query: str) -> bool:
        """验证SQL是否为查询语句

        Args:
            query: SQL查询语句

        Returns:
            bool: 如果是SELECT语句返回True，否则返回False
        """
        # 移除注释和多余的空白字符
        clean_query = re.sub(r'--.*$|/\*[\s\S]*?\*/', '', query, flags=re.MULTILINE).strip()
        # 检查是否以SELECT开头（忽略大小写）
        return bool(re.match(r'^SELECT\s', clean_query, re.IGNORECASE))
